feat: [grpc] Add retry logic when certificate mismatch for existing credentials & Agent Identity workloads - #18019
Conversation
feat: Add passphrase handling to client cert callback
feat: Add cert rotation handling support
chore: Add passphrase in _mtls_helper for requests
…rase chore: Modify _mtls_helper call to include additional variable passphrase
Updated mock return values in test cases to include None for additional parameters.
Updated mock callback to return an additional None value.
Updated mock call_client_cert_callback to include a None value in the return tuple.
chore: Add unit tests for cert rotation handling for grpc
chore: Add tests for grpc cert roatation handling changes
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces mTLS certificate rotation and automatic retry capabilities for gRPC transport by implementing a refreshing channel wrapper and call interceptors. It also updates helper utilities to support key passphrases. The review feedback suggests handling potential CancelledError exceptions gracefully when checking futures to prevent thread crashes, and recommends removing several redundant getattr checks for attributes that are guaranteed to be initialized in their respective class constructors.
fix: Refactor gRPC call handling and state management
|
The github status checks are stuck. You could try closing and re-opening this PR to re-trigger them |
cdd91ce to
524247b
Compare
Signed-off-by: Radhika Agrawal <agrawalradhika@google.com>
…mments and fix doctrings chore: Update Refactor mTLS gRPC client interceptor logic based on comments and fix doctrings
…ype hint for passphrase. chore: Refactor error handling for certificate retrieval and update type hint for passphrase.
fix: Refactor deadline error handling in grpc.py
Signed-off-by: Radhika Agrawal <agrawalradhika@google.com>
daniel-sanche
left a comment
daniel-sanche
left a comment
There was a problem hiding this comment.
I left some comments after a first pass, but I need to spend more time reviewing some of these wrappers.
First though, are you sure this is the right layer? I don't have much context on this change, but IIRC, gapic libraries use a different channel type. So I just want to make sure this is accomplishing what you want it to
| ) | ||
|
|
||
| return grpc.secure_channel(target, composite_credentials, **kwargs) | ||
| is_recreation = kwargs.pop("_is_recreation", False) |
There was a problem hiding this comment.
Is there a way we can add this as an advertised kwarg? Maybe call it something like auto_refresh? Having these secret parameters feels like an antipattern
| ) | ||
|
|
||
| # If SSL credentials are not explicitly set, try client_cert_callback and ADC. | ||
| cached_cert = None |
There was a problem hiding this comment.
if feasible to add type hinting here, that would be helpful, so we know what to expect for this variable later
|
|
||
| wrapper = _MTLSRefreshingChannel(target, factory_args, channel, cached_cert) | ||
|
|
||
| interceptor._wrapper = wrapper |
There was a problem hiding this comment.
can we make a method to set this instead of manipulating internal state?
Or pass in the wrapper when creating the interceptor?
| "request": request, | ||
| "target": target, | ||
| "ssl_credentials": None, | ||
| "client_cert_callback": client_cert_callback, |
There was a problem hiding this comment.
why is the callback being passed in? It looks like it is set to None later?
| "client_cert_callback": client_cert_callback, | ||
| "_is_recreation": True, # Hidden flag to stop recursion | ||
| **kwargs, | ||
| } |
There was a problem hiding this comment.
Instead of passing through all these arguments, it could be better to create a create_channel_fn partial here, and pass that in instead (we did something similar for bigtable)
| old_channel.unsubscribe(callback) | ||
| except Exception: | ||
| pass | ||
| self._channel.subscribe(callback) |
There was a problem hiding this comment.
does the old_channel need to be closed?
| return self._call.time_remaining() | ||
|
|
||
| def add_callback(self, callback): | ||
| self._call.add_callback(callback) |
There was a problem hiding this comment.
I'm seeing some repeated boilerplate between this and _RetryableUnaryResponseFuture. Would it help to make a generic wrapper, that both can inherit from?
| ) | ||
|
|
||
|
|
||
| class _DeadlineExceededError(grpc.RpcError, grpc.Call): |
There was a problem hiding this comment.
Why does this need to be a grpc.Call?
| return self._is_mtls | ||
|
|
||
|
|
||
| class _MTLSCallInterceptor( |
There was a problem hiding this comment.
It might be worth creating a new file for these new classes? There are a bunch of them that seem pretty tightly coupled, without much integration with other classes in this file
4 participants
feat: [grpc] Add retry logic when certificate mismatch for existing credentials & Agent Identity workloads
This PR introduces mTLS certificate rotation and encrypted private key passphrase handling to the gRPC transport. It implements an interceptor and a refreshing channel to automatically reload client certificates upon authentication failures when certificates rotate. It also updates tests and mock return values across all transports.
See go/grpc-cert-rotation-in-pythonsdk-for-x509 for details.